package com.health.trans.pojo;

import cn.hutool.extra.spring.SpringUtil;
import com.health.trans.processor.listener.MyItemReadListener;
import com.health.trans.processor.listener.StepItemWriteListener;
import com.health.trans.util.JobUtil;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.springframework.batch.core.Job;
import org.springframework.batch.core.job.builder.JobBuilder;
import org.springframework.batch.core.listener.JobExecutionListenerSupport;
import org.springframework.batch.core.step.builder.AbstractTaskletStepBuilder;
import org.springframework.batch.core.step.builder.StepBuilder;
import org.springframework.batch.core.step.tasklet.TaskletStep;
import org.springframework.batch.item.ItemProcessor;
import org.springframework.batch.item.ItemReader;
import org.springframework.batch.item.ItemWriter;

import java.util.function.Function;

import static com.health.trans.service.BatchTransService.myTaskExecutor;

/**
 * @author nqlz
 */
@Data
@NoArgsConstructor
public class JobInit<T> {
    private ItemReader<T> reader;
    private ItemWriter<T> writer;
    private ItemProcessor itemProcessor;
    private int chunkSize;
    private String jobName;
    private JobExecutionListenerSupport endListener;
    private Job job;
    private Boolean isParallel;
    private Integer parallelNum = 5;

    public JobInit(int chunkSize, String jobName) {
        this.chunkSize = chunkSize;
        this.jobName = jobName;
        this.stepBuilder = JobUtil.stepBuilderFactory.get(jobName);
        this.jobBuilder = JobUtil.jobBuilderFactory.get(jobName);
    }

    private TaskletStep step;
    private StepBuilder stepBuilder;
    private JobBuilder jobBuilder;
    private AbstractTaskletStepBuilder<?> simpleStepBuilder;

    public void initStep(Function<StepBuilder, StepBuilder> stepSupplier) {
        stepBuilder = stepSupplier.apply(stepBuilder);
        initStep();
    }

    public void initSimpleStep() {
        MyItemReadListener readListener = SpringUtil.getBean(MyItemReadListener.class);
        StepItemWriteListener writeListener = SpringUtil.getBean(StepItemWriteListener.class);
        simpleStepBuilder = !isParallel ? stepBuilder
                .<T, T>chunk(chunkSize)
                .reader(reader)
                .listener(readListener)
                .processor(itemProcessor)
                .writer(this.writer).listener(writeListener)
                : stepBuilder
                .<T, T>chunk(chunkSize)
                .reader(reader)
                .listener(readListener)
                .processor(itemProcessor)
                .writer(this.writer)
                .listener(writeListener)
                .taskExecutor(myTaskExecutor)
                .throttleLimit(parallelNum);
    }

    public void initStep() {
        if (simpleStepBuilder == null) {
            initSimpleStep();
        }
        step = simpleStepBuilder.build();
    }

    public void initJob() {
        if (step == null) {
            initStep();
        }
        job = jobBuilder.listener(endListener)
                .flow(step)
                .end().build();
    }

    public void initJob(Function<JobBuilder, JobBuilder> jobBuilderConsumer) {
        if (step == null) {
            initStep();
        }
        jobBuilder = jobBuilderConsumer.apply(jobBuilder);
        initJob();
    }

    public Job getJob() {
        if (job == null) {
            initJob();
        }
        return job;
    }

}
